;; Multiarray implementation for Unidata array class

;; This module requires the netCDF Java library (version 4)
;; on the classpath.

;; by Konrad Hinsen
;; last updated February 8, 2010

;; Copyright (c) Konrad Hinsen, 2009-2010. All rights reserved.  The use
;; and distribution terms for this software are covered by the Eclipse
;; Public License 1.0 (http://opensource.org/licenses/eclipse-1.0.php)
;; which can be found in the file epl-v10.html at the root of this
;; distribution.  By using this software in any fashion, you are
;; agreeing to be bound by the terms of this license.  You must not
;; remove this notice, or any other, from this software.

(use 'nstools.ns)
(ns+ multiarray.unidata
  (:clone multiarray.utility)
  (:use multiarray.protocol))

;;
;; Array constructors
;;

(defn zeros
  ([shape]
     (zeros shape Integer/TYPE))
  ([shape element-class]
     (ucar.ma2.Array/factory element-class (int-array shape))))

(defn ma-range
  ([end]
     (ma-range 0 end 1))
  ([start end]
     (ma-range start end 1))
  ([start end step]
     (assert (every? number? [start end step]))
     (let [class-rank {Byte 0 Short 0 Integer 0
		       clojure.lang.Ratio 1
		       Float 2 Double 2}
	   element-table {0 Integer/TYPE 1 Object 2 Double/TYPE}
	   element-class (element-table
			  (apply max (map (comp class-rank class)
					  [start end step])))
	   n (if (identical? element-class Integer/TYPE)
	       (count-range start end step)
	       (count (range start end step)))
	   ma (ucar.ma2.Array/factory element-class (int-array [n]))
	   iter (.getIndexIterator ma)]
       (doseq [n (range start end step)]
	 (.setObjectNext iter n))
       ma)))

;;
;; Implementation of multiarray.protocol
;;

(defn- unidata-sample
  [#^ucar.ma2.Array ma #^Integer dim
   #^Integer start #^Integer end #^Integer step]
  (let [rank (.getRank ma)
	ranges (new java.util.ArrayList rank)]
    (int-loop [i 0 rank]
      (.add ranges (if (= i dim)
		     (new ucar.ma2.Range start (dec end) step)
		     nil)))
    (.sectionNoReduce ma ranges)))

; try the array-argument form of sectionNoReduce to see if it handles
; empty result arrays correctly...
; .. but it doesn't.
(defn ud-sample
  [#^ucar.ma2.Array ma #^Integer dim
   #^Integer start #^Integer end #^Integer step]
  (let [rank     (.getRank ma)
	shape    (.getShape ma)
	origin   (int-array rank)
	subshape (int-array rank)
	strides  (int-array rank)]
    (int-loop [i 0 rank]
      (if (= i dim)
        (do (aset-int origin i start)
            (aset-int subshape i (count-range start end step))
            (aset-int strides i step))
        (do (aset-int origin i 0)
            (aset-int subshape i (aget shape i))
            (aset-int strides i 1))))
    (.sectionNoReduce ma origin subshape strides)))

(extend ucar.ma2.Array
  MultiArrayProtocol
  (merge multiarray-method-defaults
	 {:shape         (fn [#^ucar.ma2.Array ma] (vec (.getShape ma)))
	  :rank          (fn [#^ucar.ma2.Array ma] (.getRank ma))
	  :sample        unidata-sample})
  MultiArrayLowLevelProtocol
  (merge multiarray-lowlevel-method-defaults
	 {:element-class (fn [#^ucar.ma2.Array ma]
			   (.getElementType ma))
	  :element-array (fn [#^ucar.ma2.Array ma]
			   (.get1DJavaArray ma (.getElementType ma)))}))
